/*
Copyright (C) 2011 by the PSVN Research Group, University of Alberta
*/

#include <vector>
#include <dlfcn.h>
#include <inttypes.h>
#include <sys/time.h>
#include "astar_data.hpp"
#include "reservoir.hpp"
#include "psvn_game_so.h"
#include "timer.hpp"

extern "C" {
#include "pdb.h"
}

using namespace std;

char* print_bytes(size_t bytes)
{
    static char str[128];
    if (bytes < 1024)
        sprintf(str, "%lub", bytes);
    else if (bytes < 1024*1024)
        sprintf(str, "%lukb", bytes / 1024);
    else if (bytes < 1024*1024*1024)
        sprintf(str, "%lumb", bytes / (1024*1024));
    else
        sprintf(str, "%.2lfgb", (double)bytes / (1024*1024*1024));
    return str;
}

void get_state_space_name(const char* in, char* out)
{
    const char* str = strrchr(in, '/');
    if (*str == 0)
        str = in;
    str++;
    while (*str && (isalpha(*str) || isdigit(*str) || *str=='_')) {
        *out++ = *str++;
    }
    *out = 0;
}

class ChildGenerator
{
public:
    ChildGenerator(const compiled_game_so_t* game, const void* state);
    
    bool Generate(void* child);

    int Cost() const;

    const char* RuleName() const;

private:
    const compiled_game_so_t* m_game;
    const void* m_state;
    so_func_ptr m_iter;
    int m_rule_used;
};

ChildGenerator::ChildGenerator(const compiled_game_so_t* game, const void* state)
    : m_game(game),
      m_state(state)
{
    init_so_fwd_iter(m_iter, m_game);
}

bool ChildGenerator::Generate(void *child)
{
    if ((m_rule_used = next_so_fwd_iter(m_iter, m_state)) >= 0) {
        apply_so_fwd_rule(m_rule_used, m_state, child, m_game);
        return true;
    }
    return false;
}

int ChildGenerator::Cost() const
{
    return so_fwd_rule_cost(m_rule_used, m_game);
}
       
const char* ChildGenerator::RuleName() const 
{
    return m_game->fwd_rule_names[m_rule_used];
}

int Heuristic(const void* state, const pdb_t* pdb, const compiled_game_so_t* game)
{
    if (is_so_goal(state, game))
        return 0;
    int epsilon = cost_of_cheapest_so_fwd_rule(game);
    int h = 0;
    if (pdb != NULL)
        h = pdb_lookup(pdb, state, game);
    return std::max(epsilon, h);
}

int AStar(const compiled_game_so_t* game, const pdb_t* pdb, const void* start_state,
          int64_t *nodes_expanded, int64_t *nodes_stored, int64_t *memory_used)
{
    if (is_so_goal(start_state, game))
        return 0;

    Reservoir res(16);
    void* child = res.Get(game->state_size);
    AStarData data;
    data.AddData(hash_so_state(start_state, game), StateData(start_state), 0, 0);

    *nodes_expanded = 0;
    while (!data.OpenListEmpty()) {
        int64_t current_hash;
        StateData current_data;
        data.Pop(current_hash, current_data);
        (*nodes_expanded)++;
        ChildGenerator generator(game, current_data.State());
        while (generator.Generate(child)) {
            const int64_t child_hash = hash_so_state(child, game);
            const int child_g = current_data.G() + generator.Cost();
            // relax previously seen child
            if (data.Relax(child_hash, child_g, current_data.State()))
                continue;
            // add new child
            void* new_child = res.Get(game->state_size);
            copy_so_state(new_child, child, game);
            const int child_h = Heuristic(child, pdb, game);
            data.AddData(child_hash, StateData(new_child, current_data.State(),
                                               child_g, child_h),
                         child_g + child_h, child_g);
        }
        if (is_so_goal(current_data.State(), game)) {
            *nodes_stored = data.NumStored();
            *memory_used = data.Memory();
            return current_data.G();
        }
    }
    return INT_MAX;
}

int main( int argc, char **argv )
{
    void *so_handle, *state;
    int64_t expanded, total_expanded, stored, total_stored, memory_used;
    int trials, d;
    compiled_game_so_t *game;
    void* pdb;
    char line[4096];

    /* read the game */
    if( argc < 2 ) {
        fprintf( stderr, "requires a shared game object\n" );
        exit( EXIT_FAILURE );
    }
    so_handle = dlopen( argv[1], RTLD_LAZY );
    if( so_handle == NULL ) {
        fprintf( stderr, "could not open shared object %s: %s\n",
                 argv[ 1 ], dlerror() );
        exit( EXIT_FAILURE );
    }
    char name[1024];
    get_state_space_name(argv[1], name);
    game = (compiled_game_so_t*)dlsym(so_handle, name);
    if (game == NULL) {
        game = (compiled_game_so_t*)dlsym( so_handle, "psvn_state_space" );
        if( game == NULL ) {
            fprintf( stderr, "could not read game object from %s\n", argv[1]);
            exit( EXIT_FAILURE );
        }
    }

    /* read the pdb */
    if( argc > 2 ) {
        pdb = read_pdb( argv[2], game );
        if (pdb == NULL)
            return EXIT_FAILURE;
    } else {
        pdb = NULL;
    }
    
    state = new_so_state( game );
    total_expanded = 0;
    total_stored = 0;
    Timer total_timer;
    for( trials = 0;
         fgets( line, 4096, stdin ) != NULL
             && read_so_state( line, state, game ) > 0;
         ++trials ) 
    {
        printf( "problem %d: ", trials + 1 );
        print_so_state( stdout, state, game );
        printf( "\n" );
        
        Timer instance_timer;
        d = AStar( game, (pdb_t*)pdb, state, &expanded, &stored, &memory_used);
        assert(d != INT_MAX);

        printf("elapsed %0.3lf cost %d expanded %lu stored %lu memory %s\n", 
               instance_timer.Elapsed(), d, expanded, stored, 
               print_bytes(memory_used));
        total_expanded += expanded;
        total_stored += stored;
    }
    destroy_so_state(state);
    if( pdb != NULL )
        destroy_pdb( (pdb_t*)pdb, game );
    double elapsed = total_timer.Elapsed();
    printf("==================================\n");
    printf("expanded %lu (%.3gm) %.2fm/s\n", total_expanded,
           (total_expanded / 1000000.0),
           (total_expanded / 1000000.0) / elapsed);
    printf("stored   %lu (%.3gm) %.2fm/s\n", total_stored,
           (total_stored / 1000000.0),
           (total_stored / 1000000.0) / elapsed);
    printf("elapsed  %.03lfs\n", elapsed);
    return 0;
}
